/*
 *  Copyright (C) 2022  贺龙宇
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program. If not, see <http://www.gnu.org/licenses/>.
 */
#include "RuneDetector.h"

// #define COLOR 1 //0红色 1蓝色

RuneDetector::RuneDetector(){

    // rad=57.2957795;
    // pi2=3.14159265359*2;
    // frame_interval_s=0.02;

    // predict_time = 0.3; //预测超前时间，单位s
    
    // blue_thresh=180;
    // red_thresh=100;

    // blue_color_thresh=0;
    // red_color_thresh=0;
    
    
    // r_max_ratio=1.7;
    // r_min_area=200;
    // r_max_area=900;

    // leaf_min_area=4500;
    // leaf_max_area=8500;
    // leaf_min_ratio=1.5;
    // leaf_max_ratio=2.7;
    
    // armor_min_area=320;
    // armor_max_area=1500;
    // armor_min_ratio=1.0;
    // armor_max_ratio=2.6;

    // R2armor_min_dist=100;
    // R2armor_max_dist=150;
    // // arm_r=75;
    
    // frame_interval_s=0.1;


    isLost=true;

/*以下为计算测速间隔帧数的代码，用定帧数据（NUC 200 G 30 C 100）乘上测速间隔时间*/
#ifdef MANIFOLD_G
    // frame_interval_s = 0.033; //帧间隔，单位s
    test_spd_frame=30*runeParam.frame_interval_s;
#endif
#ifdef MANIFOLD_C
    // frame_interval_s=0.01;
    test_spd_frame=100*runeParam.frame_interval_s;
#endif
#ifdef NUC
    // frame_interval_s=0.005;
    test_spd_frame=100*runeParam.frame_interval_s;
#endif
    
#ifdef DEBUG
    show_graph=cv::Mat::zeros(cv::Size(640,480),CV_8UC3);
#endif

    /*装完B就跑*/
    

}

// std::string stCurrentTime;
// cv::VideoWriter video_writer("1.avi",cv::VideoWriter::fourcc('M','J','P','G'), 30, cv::Size(640, 480), true);

bool RuneDetector::DetectRune(const cv::Mat &src, cv::Point3f &target){
    R_tag_found=false;
    armor_found=false;

    cv::RotatedRect target_rect,R_tag;
    GetTarget(src,target_rect,R_tag);

    
    /*如果第一次没识别到，则返回上一时刻R标的位置。
    如果上一时刻也没识别到R标，则认为没有识别到目标，返回（0，0，-1）*/

    static cv::Point2f last_R_center;
    if(R_tag_found==false || armor_found==false){
        isLost=true;
        if(last_R_center.x>0.0){
            SolvePosition(last_R_center,target);
            last_R_center.x=-1;
            return false;
        }
        return false;
    }    

    last_R_center=R_tag.center;
    


    // if(nullptr==&target_rect) return;
    // Init(R_tag,hit);
    
    /*预测击打目标*/

    cv::RotatedRect predict_rect;

    
    /*小符算法*/
    if(mcu_data.state==State::RUNE_SMALL)
        SolveAndPredict_SMALL(target_rect,R_tag,predict_rect);

    /*大符算法*/
    else if(mcu_data.state==State::RUNE_BIG)
        SolveAndPredict_BIG(target_rect,R_tag,predict_rect);

    /*将识别效果保存到视频*/
    // video_writer.write(show_final);


    predict_rect=Estimation(R_tag,predict_rect);

    /*位置解算*/

    SolvePosition(predict_rect.center,target);
    return true;
}

void RuneDetector::GetTarget(const cv::Mat &src, cv::RotatedRect &target_rect, cv::RotatedRect &R_tag){
    
    /*采用颜色二值化方案*/
    /*****************************************************************************/
    /*采用颜色二值化处理的过程*/
    cv::Mat distill;
    if(runeParam.imgproc_mode==RUNE_COLOR){
        cv::Mat channels[3];
        cv::split(src,channels);
        if(mcu_data.color == EnemyColor::BLUE)//敌方是蓝色，因此能量机关识别红色
            distill=channels[2]-channels[0];
        else 
            distill=channels[0]-channels[2];
    }
    /*******************************************************************************/
    /* 采用亮度二值化方案 */
    /*****************************************************************************/
    else{
        cv::Mat channels[3];
        cv::Mat brightness_src;
        cv::cvtColor(src,brightness_src,cv::COLOR_BGR2HSV);
        cv::split(src,channels);
        distill=channels[2];
    }
    /******************************************************************************/

    cv::Mat bin_img;
    cv::threshold(distill,bin_img,(bool)mcu_data.color?runeParam.red_thresh:runeParam.blue_thresh,255,cv::THRESH_BINARY);

    // cv::erode(bin_img,bin_img,cv::Mat(),cv::Point(-1,-1),1);
    cv::dilate(bin_img,bin_img,cv::Mat(),cv::Point(-1,-1),runeParam.ed_times);
    cv::erode(bin_img,bin_img,cv::Mat(),cv::Point(-1,-1),runeParam.ed_times);
#ifdef DEBUG
    // double s;
    // double min_value,max_value;
    // cv::Point min_loc,max_loc;
    // cv::minMaxLoc(color_distill.reshape(1),&min_value,&max_value,&min_loc,&max_loc);
    // cv::putText(color_distill,std::to_string(max_value),max_loc,0,0.5,cv::Scalar(255,255,255),2,16);
    // cv::imshow("color_distill",color_distill);
    // cv::imshow("distill",distill);
    // cv::imshow("binary",bin_img);
#endif
    
    
    std::vector<std::vector<cv::Point>>  contours;

    std::vector<cv::Vec4i> hierarchy;
    cv::findContours(bin_img,contours,hierarchy,cv::RETR_TREE,cv::CHAIN_APPROX_SIMPLE);


#ifdef DEBUG

    show_contours=cv::Mat::zeros(cv::Size(src.cols,src.rows),CV_8UC3);
    show_area=src.clone();
    show_final=src.clone();
    show_ratio=src.clone();
    show_inside_count=src.clone();
    

    for(int i=0;i<contours.size();i++){
        cv::RotatedRect rect=cv::minAreaRect(contours[i]);

        float length=std::max(rect.size.height,rect.size.width);
        float width =std::min(rect.size.height,rect.size.width);

        float ratio=length/width;
        // if(ratio>3) continue;
        cv::putText(show_area,std::to_string((int)rect.size.area()),rect.center,0,0.5,cv::Scalar(0,255,0),1,16);
        cv::putText(show_ratio,std::to_string(ratio),rect.center,0,0.3,cv::Scalar(0,255,0),1,16);
    }

    /*画轮廓*/
    for(int i=0;i<contours.size();i++){
        for(int j=0;j<contours[i].size()-1;j++){
            cv::line(show_contours,contours[i][j],contours[i][j+1],cv::Scalar(255,255,255),1,16);
        }
            cv::line(show_contours,contours[i][contours[i].size()-1],contours[i][0],cv::Scalar(255,255,255),1,16);
    }
#endif

    /*寻找待击打的装甲板*/

    for(int i=0;i<contours.size();i++){

        cv::RotatedRect rect = cv::minAreaRect(contours[i]);

        /*如果扇叶的面积、比例不符合要求，则直接识别下一个轮廓*/
        float width=std::min(rect.size.width,rect.size.height);
        float length=std::max(rect.size.width,rect.size.height);
        float ratio=length/width;
        
        if(rect.size.area()<runeParam.leaf_min_area || rect.size.area()>runeParam.leaf_max_area
           || ratio<runeParam.leaf_min_ratio || ratio>runeParam.leaf_max_ratio){
               continue;
           }

        /*找到扇叶后，判断该扇叶是否有内嵌轮廓，并且该轮廓是否唯一*/
        if(hierarchy[i][2]>0 && hierarchy[hierarchy[i][2]][0]<0 && hierarchy[hierarchy[i][2]][1]<0){
            // std::cout<<hierarchy[i][2]<<std::endl;
            target_rect=cv::minAreaRect(contours[hierarchy[i][2]]);
            armor_found=true;
            break;
        }
    }
#ifdef DEBUG
    cv::ellipse(show_final,target_rect,cv::Scalar(0,255,0),8,16);
#endif

    /*寻找R标,寻找长宽比接近1的、大小适中,且没有父轮廓的装甲板*/
    
    for(int i=0;i<contours.size();i++){
        cv::RotatedRect rect=cv::minAreaRect(contours[i]);
        float length=std::max(rect.size.height,rect.size.width);
        float width =std::min(rect.size.height,rect.size.width);

        float ratio=length/width;

        if(ratio<runeParam.r_max_ratio && rect.size.area()>runeParam.r_min_area  
            && rect.size.area()<runeParam.r_max_area && hierarchy[i][3]<0){
            /*拿轮廓上其中一个点的颜色进行判断，是否为目标颜色*/
            // if(color_distill.at<uchar>(contours[i][0])<((bool)mcu_data.color?red_color_thresh:blue_color_thresh))
                // continue;

            /*若已识别到装甲板，则判断R标与装甲板的距离是否满足条件。若不满足，则继续*/

            if(armor_found){
                float dist_x=rect.center.x-target_rect.center.x;
                float dist_y=rect.center.y-target_rect.center.y;
                float distance=std::sqrt(dist_x*dist_x+dist_y*dist_y);
                
                // std::cout<<distance<<std::endl;

                if(distance<runeParam.R2armor_min_dist || distance>runeParam.R2armor_max_dist)
                    continue;
            }

            R_tag=rect;
            R_tag_found=true;
#ifdef DEBUG
            cv::ellipse(show_final,rect,cv::Scalar(0,0,255),2,16);
#endif
            break;
        }
    }
#ifdef DEBUG
    cv::putText(show_final,R_tag_found?"R Found":"R Not Found",cv::Point(50,150),
            0,1,R_tag_found?cv::Scalar(0,255,0):cv::Scalar(0,0,255),2,16);
    cv::putText(show_final,armor_found?"Armor Found":"Armor Not Found",cv::Point(50,200),
            0,1,armor_found?cv::Scalar(0,255,0):cv::Scalar(0,0,255),2,16);
    cv::imshow("contours",show_contours);
    cv::imshow("show_area",show_area);
    cv::imshow("show_ratio",show_ratio);
    cv::imshow("show_final",show_final);
    // cv::imshow("show_inside_count",show_inside_count);
#endif
}

void RuneDetector::SolveAndPredict_SMALL(const cv::RotatedRect &target_rect, 
                     const cv::RotatedRect &R_tag, 
                     cv::RotatedRect &predict_rect)
{
    // static int exe_cnt=0;

    float delta_y=R_tag.center.y-target_rect.center.y;
    float delta_x=target_rect.center.x-R_tag.center.x;
    float arm_r=std::sqrt(delta_x*delta_x+delta_y*delta_y);


    // std::cout<<arm_r<<std::endl;
    
    
    float angle=std::atan2(delta_y,delta_x);

    // angle*=rad;

    /*计算角速度*/
    /*优化：取0.1s为测速区间，对应G，C，NUC为3帧、10帧、10帧*/
    static float last_angle;
    static float velocity=0.0;
    float next_angle;
    static float correct_val; /*用于处理阶跃点的情况*/
    static int count=0;  /*用于计数，判断进行了多少帧*/
    static float start_angle;
    
    /*若上一时刻角度跟这一时刻相差5弧度（将近一个圆），说明发生了突变，根据情况更新correct_val*/
    if(std::abs(angle-last_angle)>5){
        /*从大到小突变，说明角度是往增大方向的，设为2pi*/
        if(last_angle>0){
            correct_val+=runeParam.pi2;
        }
        /*从大到小突变，说明角度往减小方向，设为-2pi*/
        else {
            correct_val+=(-runeParam.pi2);
        }
    }

    angle+=correct_val;

    /*若count=test_spd_frame,则计算一次速度，并初始化count和correct_val*/
    /*此处需要做一个判断。如果中途发生了目标切换，则不能更新速度。否则云台会直接跟丢*/

    /*如果出现了掉帧的情况，则重置count*/

    if(isLost){
        count=0;
        isLost=false;
    }

    if(count==test_spd_frame){
        if(std::abs(angle-start_angle)<0.5) /*角度变化在0.5弧度内，才更新速度值*/
            velocity=(angle-start_angle)/runeParam.frame_interval_s;
        count=0;
        correct_val=0;
    }
    /*若count=0，则将当前angle设为start_angle*/
    else if(count==0){
        start_angle=angle;
        count++;
    }

    
    /*否则直接令count++*/
    else{
        count++;
    }

        next_angle=angle+velocity*runeParam.predict_time; /*预测下一时刻角度位置*/

        /*计算预测矩形*/
        predict_rect=target_rect;
        predict_rect.center.x=R_tag.center.x+arm_r*std::cos(next_angle);
        predict_rect.center.y=R_tag.center.y-arm_r*std::sin(next_angle);
#ifdef DEBUG
        cv::circle(show_final,cv::Point(R_tag.center.x+arm_r*std::cos(next_angle)
                    ,R_tag.center.y-arm_r*std::sin(next_angle)),4,cv::Scalar(255,255,255),4,16);
#endif
    while(std::abs(angle)>CV_PI){
        if(angle>0) angle-=runeParam.pi2;
        else angle+=runeParam.pi2;
    }
    last_angle=angle;

    /*计算角加速度*/
    // static float last_velocity;
    // static float accelerate;
    // if(exe_cnt>1) accelerate=velocity-last_velocity;
    // last_velocity=velocity;
    
#ifdef DEBUG
    cv::putText(show_final,"pos: "+std::to_string(angle),cv::Point(50,50),0,1,cv::Scalar(0,255,0),2,16);
    cv::putText(show_final,"spd: "+std::to_string(velocity),cv::Point(50,100),0,1,cv::Scalar(0,255,0),2,16);
    cv::putText(show_final,"radius: "+std::to_string(arm_r),cv::Point(50,350),0,1,cv::Scalar(0,255,0),2,16);
    // cv::putText(show_final,std::to_string(accelerate),cv::Point(50,150),0,1,cv::Scalar(0,255,0),2,16);

    cv::imshow("show_final",show_final);
#endif

    // exe_cnt++;
}

void RuneDetector::SolveAndPredict_BIG(const cv::RotatedRect &target_rect, 
                     const cv::RotatedRect &R_tag, 
                     cv::RotatedRect &predict_rect)
{
    // static int exe_cnt=0;

    float delta_y=R_tag.center.y-target_rect.center.y;
    float delta_x=target_rect.center.x-R_tag.center.x;
    float arm_r=std::sqrt(delta_x*delta_x+delta_y*delta_y);


    // std::cout<<arm_r<<std::endl;
    
    
    float angle=std::atan2(delta_y,delta_x);

    // angle*=rad;

    /*计算角速度*/
    /*优化：取0.1s为测速区间，对应G，C，NUC为3帧、10帧、20帧*/
    static float last_angle;
    static float velocity=0.0;
    float next_angle;
    static float correct_val; /*用于处理阶跃点的情况*/
    static int count=0;  /*用于计数，判断进行了多少帧*/
    static float start_angle;

    /*计算角加速度，此时不需要考虑阶跃的情况*/
    static float last_velocity=2022.0514;
    static float accelerate=0.0;
    
    
    /*若上一时刻角度跟这一时刻相差5弧度（将近一个圆），说明发生了突变，根据情况更新correct_val*/
    if(std::abs(angle-last_angle)>5){
        /*从大到小突变，说明角度是往增大方向的，设为2pi*/
        if(last_angle>0){
            correct_val+=runeParam.pi2;
        }
        /*从大到小突变，说明角度往减小方向，设为-2pi*/
        else {
            correct_val+=(-runeParam.pi2);
        }
    }

    angle+=correct_val;

    /*若count=test_spd_frame,则计算一次速度，并初始化count和correct_val*/
    /*此处需要做一个判断。如果中途发生了目标切换，则不能更新速度。否则云台会直接跟丢*/
    /*补充：对于大符，将帧数减半，提高一倍的更新速度*/

    if(isLost){
        count=0;
        isLost=false;
    }

    if(count==test_spd_frame){
        if(std::abs(angle-start_angle)<0.5) /*角度变化在0.5弧度内，才更新速度值和加速度值*/
        {
            velocity=(angle-start_angle)/(runeParam.frame_interval_s);

            if(last_velocity<1024)
                accelerate=(velocity-last_velocity)/runeParam.frame_interval_s;
            
            // last_velocity=velocity;
        }
        count=0;
        correct_val=0;
    }
    /*若count=0，则将当前angle设为start_angle*/
    else if(count==0){
        start_angle=angle;
        last_velocity=velocity;
        count++;
    }

    
    /*否则直接令count++*/
    else{
        count++;
    }

        next_angle=angle+velocity*runeParam.predict_time+0.5*accelerate*runeParam.predict_time*runeParam.predict_time; /*预测下一时刻角度位置*/

        /*计算预测矩形*/
        predict_rect=target_rect;
        predict_rect.center.x=R_tag.center.x+arm_r*std::cos(next_angle);
        predict_rect.center.y=R_tag.center.y-arm_r*std::sin(next_angle);
#ifdef DEBUG
        cv::circle(show_final,cv::Point(R_tag.center.x+arm_r*std::cos(next_angle)
                    ,R_tag.center.y-arm_r*std::sin(next_angle)),4,cv::Scalar(255,255,255),4,16);
#endif
    while(std::abs(angle)>CV_PI){
        if(angle>0) angle-=runeParam.pi2;
        else angle+=runeParam.pi2;
    }
    last_angle=angle;

    /*计算角加速度*/
    // static float last_velocity;
    // static float accelerate;
    // if(exe_cnt>1) accelerate=velocity-last_velocity;
    // last_velocity=velocity;
    
#ifdef DEBUG
    cv::putText(show_final,"pos: "+std::to_string(angle),cv::Point(50,50),0,1,cv::Scalar(0,255,0),2,16);
    cv::putText(show_final,"spd: "+std::to_string(velocity),cv::Point(50,100),0,1,cv::Scalar(0,255,0),2,16);
    cv::putText(show_final,"acc: "+std::to_string(accelerate),cv::Point(50,300),0,1,cv::Scalar(0,255,0),2,16);
    cv::putText(show_final,"radius: "+std::to_string(arm_r),cv::Point(50,350),0,1,cv::Scalar(0,255,0),2,16);
    // cv::putText(show_final,std::to_string(accelerate),cv::Point(50,150),0,1,cv::Scalar(0,255,0),2,16);

    cv::imshow("show_final",show_final);
#endif

    // exe_cnt++;
}

cv::RotatedRect RuneDetector::Estimation(const cv::RotatedRect &R_tag, const cv::RotatedRect &predict_rect)
{
    /*TODO 需要实际调试。在此之前可以计算，或者猜*/
    /*思路：以目标在y方向相对于R标的差值为依据进行补偿*/

    /*国赛计划使用4段法进行补偿*/
    /*从低到高依次对应runeParam中的estimation1~estimation4*/
    float delta_y=R_tag.center.y-predict_rect.center.y;
    float delta_x=predict_rect.center.x-R_tag.center.x;
    float angle=std::atan2(delta_y,delta_x);
    
    
    /*让angle保证在-pi/2到pi/2*/
    if(angle>CV_PI/2){
        angle=CV_PI-angle;
    }
    else if(angle<(-CV_PI/2)){
        angle=-CV_PI-angle;
    }

    cv::RotatedRect rect=predict_rect;

    if(angle<(-CV_PI/4))
        rect.center.y-=runeParam.estimation_1;
    else if(angle<0)
        rect.center.y-=runeParam.estimation_2;
    else if(angle<CV_PI/4)
        rect.center.y-=runeParam.estimation_3;
    else 
        rect.center.y-=runeParam.estimation_4;
    
#ifdef DEBUG
    cv::putText(show_final,"angle: "+std::to_string(angle),cv::Point(50,250),0,1,cv::Scalar(0,255,0),2,16);
    cv::circle(show_final,rect.center,4,cv::Scalar(0,255,0),2,16);
    cv::imshow("show_final",show_final);
#endif
    return rect;
}

void RuneDetector::SolvePosition(const cv::Point2f &aim_point,cv::Point3f &target)
{
    static float half_camera_width = cameraParam.resolution_width / 2;
    static float half_camera_height = cameraParam.resolution_height / 2;
    static float camera_fx = cameraParam.fx;
    static float camera_fy = cameraParam.fy;

    float yaw_offset = atan2(half_camera_width - aim_point.x, camera_fx) * runeParam.rad;
    float pitch_offset = atan2(half_camera_height - aim_point.y, camera_fy) * runeParam.rad;
    float distance = 1.000; /*给个定值即可*/

    target={yaw_offset,pitch_offset,distance};

}